import pandas as pd
import json
import re
import subprocess
from typing import List, Tuple

# === 配置 ===
INPUT_FILE = r"D:\software\Desktop\Ollama_damo\test_extract.xlsx"
OUTPUT_FILE = "fine_grained_triples.csv"
MODEL_NAME = "qwen2"
STRUCTURE_RELATIONS = {"ticket", "traffic", "tips", "comments", "season"}

def call_qwen2(prompt: str, timeout: int = 120) -> str:
    """通过 ollama run 调用 qwen2"""
    try:
        result = subprocess.run(
            ["ollama", "run", MODEL_NAME],
            input=prompt,
            text=True,
            encoding="utf-8",
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            timeout=timeout
        )
        return result.stdout.strip() if result.returncode == 0 else ""
    except Exception as e:
        print(f"[ERROR] Qwen2 调用失败: {e}")
        return ""

def extract_fine_triples(subject: str, relation: str, text: str) -> List[Tuple[str, str, str]]:
    if not text or pd.isna(text) or relation not in STRUCTURE_RELATIONS:
        return [(subject, relation, text)]

    prompt = f"""
你是一个陕西文旅知识图谱专家。请将以下关于景点“{subject}”的“{relation}”信息，拆解为多个结构化的三元组。
要求：
1. 输出为纯 JSON 列表，每项为 {{"predicate": "关系名", "object": "值"}}
2. predicate 使用英文小写+下划线（如: adult_ticket_price, bus_line, best_season）
3. object 尽量为结构化值（数字、布尔、字符串），不要大段描述
4. 若信息缺失，返回空列表 []
5. 不要任何解释、注释、markdown

示例输入（relation="ticket"）：
成人票30元；学生票15元

示例输出：
[{{"predicate": "adult_ticket_price", "object": 30}}, {{"predicate": "student_ticket_price", "object": 15}}]

现在处理：
{relation} 信息：{text}
"""
    response = call_qwen2(prompt)

    try:
        # 清理 JSON
        clean_resp = re.sub(r"```(?:json)?", "", response).strip()
        clean_resp = re.sub(r"```.*$", "", clean_resp, flags=re.DOTALL)
        triples = json.loads(clean_resp)
        if isinstance(triples, list):
            return [(subject, item["predicate"], item["object"]) for item in triples if "predicate" in item and "object" in item]
    except Exception as e:
        print(f"[JSON ERROR] {e} | 响应: {response[:150]}...")
    return [(subject, relation, text)]  # 解析失败则保留原始

def main():
    df = pd.read_excel(INPUT_FILE, sheet_name=0)
    all_triples = []

    for idx, row in df.iterrows():
        subject = str(row.get("subject", "")).strip()
        relation = str(row.get("relation", "")).strip().lower()
        obj = row.get("object", "")

        if not subject or not relation:
            continue

        fine_triples = extract_fine_triples(subject, relation, obj)
        all_triples.extend(fine_triples)

        if (idx + 1) % 10 == 0:
            print(f"✅ 已处理 {idx + 1} 行...")

    # 保存结果
    out_df = pd.DataFrame(all_triples, columns=["subject", "predicate", "object"])
    out_df.to_csv(OUTPUT_FILE, index=False, encoding="utf-8-sig")
    print(f"\n🎉 共生成 {len(all_triples)} 条三元组，已保存至 {OUTPUT_FILE}")

if __name__ == "__main__":
    main()